#include "Header.h"

void sub_140014360()
{
    HRESULT hResult;
    IMalloc* lpIMalloc = nullptr;
    IShellFolder* lpIShellFolder = nullptr;
    PIDLIST_ABSOLUTE lpItemIdList = nullptr;
    IShellFolder2* lpIShellFolder2 = nullptr;
    IEnumIDList* lpIEnumIDList = nullptr;
    PITEMID_CHILD lpItemIdChild = nullptr;

    CoInitialize(nullptr);

    do
    {
        SHGetMalloc(&lpIMalloc);

        hResult = SHGetDesktopFolder(&lpIShellFolder);
        if (FAILED(hResult)) break;

        hResult = SHGetSpecialFolderLocation(0, CSIDL_BITBUCKET, &lpItemIdList);
        if (FAILED(hResult)) break;

        hResult = lpIShellFolder->BindToObject(lpItemIdList, nullptr, IID_IShellFolder2, (void**)&lpIShellFolder2);
        if (FAILED(hResult)) break;

        hResult = lpIShellFolder2->EnumObjects(0, SHCONTF_FOLDERS | SHCONTF_NONFOLDERS | SHCONTF_INCLUDEHIDDEN, &lpIEnumIDList);
        if (FAILED(hResult)) break;

        hResult = lpIEnumIDList->Next(1, &lpItemIdChild, nullptr);
        if (hResult == S_FALSE) break;

        //
        //  TO DO
        //

    } while (FALSE);

    if (lpIEnumIDList)
        lpIEnumIDList->Release();

    if (lpIShellFolder2)
        lpIShellFolder2->Release();

    if (lpIShellFolder)
        lpIShellFolder->Release();

    if (lpIMalloc)
        lpIMalloc->Release();

    CoUninitialize();
}

void* NtfsGet80hAndRunList(IN HANDLE hFile, OUT PUINT32 lpRunListLength, OUT PFILE_RECORD_NONRESIDENT_ATTRIBUTE_HEADER lpHeader)
{
    char strBuff[0x400] = { 0 };
    BOOL bRet;
    DWORD dwRead;
    PFILE_RECORD_HEADER lpFileRecordHeader;
    UINT16 uOffset;
    PFILE_RECORD_NONRESIDENT_ATTRIBUTE_HEADER lpFrnaHeader;
    UINT16 uRunListLength;
    PUCHAR lpRunList = nullptr;

    do
    {
        if (!lpRunListLength || !lpHeader) break;

        bRet = ReadFile(hFile, strBuff, 0x400, &dwRead, nullptr);
        if (!bRet) break;

        lpFileRecordHeader = (PFILE_RECORD_HEADER)strBuff;
        uOffset = lpFileRecordHeader->OffsetToFirstAttribute;
        while (uOffset < 0x400)
        {
            lpFrnaHeader = (PFILE_RECORD_NONRESIDENT_ATTRIBUTE_HEADER)((char*)lpFileRecordHeader + uOffset);

            if (lpFrnaHeader->AttributeType == 0x80)
            {
                uRunListLength = lpFrnaHeader->Length - lpFrnaHeader->OffsetToRunList;
                lpRunList = new UCHAR[uRunListLength];
                if (!lpRunList) break;
                RtlZeroMemory(lpRunList, uRunListLength);
                RtlCopyMemory(lpRunList, (char*)lpFrnaHeader + lpFrnaHeader->OffsetToRunList, uRunListLength);
                *lpRunListLength = uRunListLength;
                RtlCopyMemory(lpHeader, lpFrnaHeader, (ULONG_PTR)((PFILE_RECORD_NONRESIDENT_ATTRIBUTE_HEADER)0)->ch);
                break;
            }

            uOffset += lpFrnaHeader->Length;

            if (!lpFrnaHeader->Length)
                uOffset -= 8;

            if (uOffset < 0) 
                break;
        }

    } while (FALSE);

    return lpRunList;
}

typedef struct _CLUSTER_STREAM_INFO
{
    UINT64 StreamTotalBytes;
    UINT64 StreamBytes;
    UINT32 StreamSize;
}CLUSTER_STREAM_INFO, * PCLUSTER_STREAM_INFO;

void GetClusterStream(
    IN void* lpClassObject,
    IN PFILE_RECORD_NONRESIDENT_ATTRIBUTE_HEADER lpFrnaHeader,
    IN PUCHAR lpRunList,
    IN UINT32 uRunListLength,
    OUT std::list<CLUSTER_STREAM_INFO>& list,
    IN PDOS_BOOT_RECORD lpBootRecord
)
{
    if (!lpFrnaHeader || !lpRunList || !uRunListLength || !lpBootRecord) return;

    UINT64 uBodySize = lpFrnaHeader->SizeOfBody;
    UINT32 uClusterBytes = lpBootRecord->BytesPerSector * lpBootRecord->SectorsPerCluster;
    BYTE uch = *lpRunList >> 4;
    BYTE uch2 = *lpRunList & 0xF;
    UINT32 uOffset = 0;
    UINT32 uStreamSize;
    PUCHAR lpCh = lpRunList;
    UINT64 uStreamBytes;
    UINT64 v26;
    UINT64 uStartOfStream;
    UINT64 v17;
    UINT32 v23 = 0;
    UINT64 v27 = 0;
    UINT64 v28[2] = { 0 };
    CLUSTER_STREAM_INFO StreamInfo = { 0 };

    if (!uch2) return;

    do
    {
        if (!uch || uOffset > uRunListLength || uch2 >= 5 || uch > 4) break;

        uStreamSize = 0;
        v26 = uBodySize;

        if (uch2)
        {
            if (lpCh)
                RtlCopyMemory(&uStreamSize, lpCh + 1, uch2);
            else
                RtlZeroMemory(&uStreamSize, uch2);
        }

        uStreamBytes = uClusterBytes * uStreamSize;
        if (!lpCh[uch + uch2 + 1])
            uStreamBytes = v26;
        v26 -= uStreamBytes;

        uStartOfStream = 0;
        if (uch2)
        {
            if (lpCh + uch2 + 1)
                RtlCopyMemory(&uStartOfStream, lpCh + uch2 + 1, uch);
            else
                RtlZeroMemory(&uStartOfStream, uch);
        }

        v17 = uStartOfStream;

        if (((1i64 << (8 * uch) >> 1) & uStartOfStream) != 0 && v23 > 0)
            v17 = uStartOfStream - (1i64 << (8 * uch));

        v27 += uClusterBytes * v17;

        StreamInfo.StreamTotalBytes = v27;
        StreamInfo.StreamBytes = uStreamBytes;
        StreamInfo.StreamSize = uStreamSize;
        list.push_back(StreamInfo);

        v23++;
        uOffset += uch2 + uch + 1;
        lpCh += uch2 + uch + 1;
        uch = *lpCh >> 4;
        uch2 = *lpCh & 0xF;

    } while ((*lpCh & 0xF) != 0);
}

BOOL NtfsGetFileNameAndTime(
    IN PFILE_RECORD_RESIDENT_ATTRIBUTE_HEADER lpHeader30H,
    IN PFILE_RECORD_ATTRIBUTE_BODY_30H lpBody30H,
    IN void* lpFileInfoClassObject,
    OUT PBOOL lpRet
)
{
    BOOL bRet = FALSE;
    PWCHAR lpWch = nullptr;

    do
    {
        if (!lpHeader30H || !lpBody30H || !lpRet) break;

        if (lpHeader30H->NonResident) break;

        *lpRet = lpBody30H->FileNameSpace == 2;

        if (!lpBody30H->FileNameLength) break;
        
        lpWch = new WCHAR[lpBody30H->FileNameLength + 1];
        if (!lpWch) break;
        RtlZeroMemory(lpWch, (lpBody30H->FileNameLength + 1) * sizeof(WCHAR));
        RtlCopyMemory(lpWch, lpBody30H->FileName, lpBody30H->FileNameLength * sizeof(WCHAR));
        printf("%ls\r\n", lpWch);

        bRet = TRUE;

    } while (FALSE);

    if (lpWch)
        delete lpWch;

    return bRet;
}

BOOL NtfsGetFileStream(
    IN void* lpClassObject,
    IN PFILE_RECORD_NONRESIDENT_ATTRIBUTE_HEADER lpHeader80H,
    IN PUCHAR lpRunList80H,
    IN UINT64 a4,
    IN void* lpClassObject2,
    IN void* lpFileInfoClassObject,
    IN UINT64 a7,
    OUT std::list<CLUSTER_STREAM_INFO>& list,
    IN PDOS_BOOT_RECORD lpBootRecord
)
{
    BOOL bRet = FALSE;
    UINT64 lpAddress;
    UINT32 uRunListLength;
    PUCHAR lpRunList = nullptr;

    do
    {
        if (lpHeader80H->Length == lpHeader80H->OffsetToAttributeName) break;

        lpAddress = 0 + a7; //  arg5 + arg7

        if (lpHeader80H->NonResident)
        {
            if (lpHeader80H->Length <= 0x40) break;

            uRunListLength = lpHeader80H->Length - lpHeader80H->OffsetToRunList;
            if (uRunListLength <= 0) break;

            lpRunList = new UCHAR[uRunListLength];
            if (!lpRunList) break;
            RtlZeroMemory(lpRunList, uRunListLength);
            RtlCopyMemory(lpRunList, lpRunList80H, uRunListLength);

            GetClusterStream(nullptr, lpHeader80H, lpRunList, uRunListLength, list, lpBootRecord);
        }
        else
        {

        }

        bRet = TRUE;

    } while (FALSE);

    if (lpRunList)
        delete[] lpRunList;

    return bRet;
}

void QuickScan()
{
    HANDLE hDrive = INVALID_HANDLE_VALUE;
    ULONG uFlags = FILE_FLAG_NO_BUFFERING | FILE_FLAG_SEQUENTIAL_SCAN | FILE_ATTRIBUTE_NORMAL | FILE_ATTRIBUTE_READONLY;
    LONG uHigh = 0;
    DWORD dwMove;
    char strBuff[0x200] = { 0 };
    BOOL bRet;
    DWORD dwRead;
    DOS_BOOT_RECORD NtfsBootRecord = { 0 };
    LARGE_INTEGER liStartOfMft = { 0 };
    UINT32 uRunListLength = 0;
    FILE_RECORD_NONRESIDENT_ATTRIBUTE_HEADER FrnaHeader = { 0 };
    PUCHAR lpRunList = nullptr;
    std::list<CLUSTER_STREAM_INFO> lstStreamInfo;
    UINT64 uClusterTotalNum;
    SYSTEM_INFO SystemInfo = { 0 };
    UINT64 v21;
    UINT64 v24;
    LARGE_INTEGER LargeInteger;
    DWORD dwMemSize;
    LPVOID lpAddress;
    UINT64 uStreamTotalSize;
    UINT64 v30;
    UINT64 uStreamRestSize;
    UINT64 v65;
    UINT32 uRecordOffset;
    PFILE_RECORD_HEADER lpRecordHeader;
    UINT64 v55;
    UINT16 uAttributeOffset;
    UINT32 uFileRecordId;
    PUCHAR lpCh;
    PFILE_RECORD_RESIDENT_ATTRIBUTE_HEADER lpFraHeader;
    UINT64 v35;
    BOOL v47;
    UINT64 v64;
    std::list<CLUSTER_STREAM_INFO> lstFileStreamInfo;
    BOOL v38;
    UINT32 v39;
    UINT32 uCnt;

    do
    {
        hDrive = CreateFile(_T("\\\\.\\E:"), GENERIC_READ, 
            FILE_SHARE_READ | FILE_SHARE_WRITE, nullptr, OPEN_EXISTING, uFlags, nullptr);
        if (hDrive == INVALID_HANDLE_VALUE) break;

        dwMove = SetFilePointer(hDrive, 0, &uHigh, FILE_BEGIN);
        if (dwMove == INVALID_SET_FILE_POINTER) break;

        bRet = ReadFile(hDrive, strBuff, sizeof(strBuff), &dwRead, nullptr);
        if (!bRet) break;

        RtlCopyMemory(&NtfsBootRecord, strBuff, sizeof(DOS_BOOT_RECORD));
        if (NtfsBootRecord.BytesPerSector != 0x200) break;

        liStartOfMft.QuadPart = NtfsBootRecord.StartClusterOfMft * NtfsBootRecord.BytesPerSector * NtfsBootRecord.SectorsPerCluster;
        dwMove = SetFilePointer(hDrive, (LONG)liStartOfMft.LowPart, &liStartOfMft.HighPart, FILE_BEGIN);
        if (dwMove == INVALID_SET_FILE_POINTER) break;

        lpRunList = (PUCHAR)NtfsGet80hAndRunList(hDrive, &uRunListLength, &FrnaHeader);
        if (!lpRunList) break;
        printf("%X\r\n", uRunListLength);

        for (unsigned int i = 0; i < uRunListLength; i++)
            printf("%02X ", ((PUCHAR)lpRunList)[i]);
        puts("");

        GetClusterStream(nullptr, &FrnaHeader, lpRunList, uRunListLength, lstStreamInfo, &NtfsBootRecord);

        uClusterTotalNum = 0;
        for (auto it = lstStreamInfo.begin(); it != lstStreamInfo.end(); ++it)
            uClusterTotalNum += (*it).StreamSize;
        printf("uClusterTotalNum: %llX\r\n", uClusterTotalNum);

        v21 = uClusterTotalNum * NtfsBootRecord.SectorsPerCluster / 2;
        printf("v21: %llX\r\n", v21);

        GetSystemInfo(&SystemInfo);
        dwMemSize = (0x12C * SystemInfo.dwAllocationGranularity) & 0xFFFFFC00;
        lpAddress = VirtualAlloc(nullptr, dwMemSize, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
        if (!lpAddress) break;

        for (auto it = lstStreamInfo.begin(); it != lstStreamInfo.end(); ++it)
        {
            v24 = (*it).StreamTotalBytes;
            LargeInteger.QuadPart = (*it).StreamTotalBytes;
            dwMove = SetFilePointer(hDrive, (LONG)LargeInteger.LowPart, &LargeInteger.HighPart, FILE_BEGIN);
            if (dwMove == INVALID_SET_FILE_POINTER) break;

            uStreamTotalSize = ((*it).StreamSize * NtfsBootRecord.SectorsPerCluster) << 9;
            uStreamRestSize = uStreamTotalSize;
            Label2:
            v30 = dwMemSize;
            if (dwMemSize > uStreamTotalSize)
                v30 = uStreamTotalSize;
            uStreamRestSize = uStreamTotalSize - v30;

            RtlZeroMemory(lpAddress, dwMemSize);
            bRet = ReadFile(hDrive, lpAddress, (DWORD)v30, &dwRead, nullptr);
            if (!bRet) break;
            v65 = v30 / 0x400;

            uRecordOffset = 0;
            uCnt = 0;
            while (TRUE)
            {
            Label:

                if (uCnt >= v65)
                {
                    v24 += v30;
                    uStreamTotalSize = uStreamRestSize;
                    if (!uStreamTotalSize) break;
                    goto Label2;
                }

                lpRecordHeader = (PFILE_RECORD_HEADER)(uRecordOffset + (ULONG_PTR)lpAddress);
                if (*((DWORD*)lpRecordHeader) != 0x454C4946) break;

                v55 = uRecordOffset + v24;
                uAttributeOffset = lpRecordHeader->OffsetToFirstAttribute;

                v35 = 0;
                v47 = TRUE;
                if (uAttributeOffset < 0x400)
                {
                    uFileRecordId = lpRecordHeader->FileRecordId;

                    while (TRUE)
                    {
                        lpCh = (PUCHAR)lpRecordHeader + uAttributeOffset;
                        lpFraHeader = (PFILE_RECORD_RESIDENT_ATTRIBUTE_HEADER)lpCh;

                        v38 = (lpFraHeader->Length & 0x3FF) == 0;
                        v39 = lpFraHeader->Length & 0x3FF;
                        lpFraHeader->Length = v39;

                        if (v38)
                        {
                            uRecordOffset += 0x400;
                            uCnt++;
                            goto Label;
                        }

                        switch (lpFraHeader->AttributeType)
                        {
                        case 0x10:
                        case 0x20:
                        case 0x40:
                        case 0x50:
                        case 0x60:
                        case 0x70:
                            uAttributeOffset += lpFraHeader->Length;
                            if (uAttributeOffset < 0x400) continue;
                            break;
                        case 0x30:
                            lpCh = (PUCHAR)lpRecordHeader + uAttributeOffset + lpFraHeader->OffsetToAttributeBody;
                            NtfsGetFileNameAndTime(lpFraHeader, (PFILE_RECORD_ATTRIBUTE_BODY_30H)lpCh, nullptr, &v47);
                            uAttributeOffset += lpFraHeader->Length;
                            if (uAttributeOffset < 0x400) continue;
                            break;
                        case 0x80:
                            v64 = uAttributeOffset + v55;
                            lpCh = (PUCHAR)lpRecordHeader + uAttributeOffset + ((PFILE_RECORD_NONRESIDENT_ATTRIBUTE_HEADER)lpFraHeader)->OffsetToRunList;
                            NtfsGetFileStream(nullptr, (PFILE_RECORD_NONRESIDENT_ATTRIBUTE_HEADER)lpFraHeader,
                                lpCh, v64, nullptr, nullptr, v55, lstFileStreamInfo, &NtfsBootRecord);
                            break;
                        default:
                            break;
                        }
                        break;
                    }
                }
                uRecordOffset += 0x400;
                uCnt++;
            }
        }

    } while (FALSE);

    if (hDrive != INVALID_HANDLE_VALUE)
        CloseHandle(hDrive);
}

int main()
{
    QuickScan();
    system("pause");
    return 0;
}